from Model.ModelTransfer import ModelFS
from Model.ModelValidate import Checker
from Model.ModelCompress import Compressor
from Model.ModelEvaluate import Appraiser
from Util.Process import INFO
from Util.SQL import update, select
from Util.Return import *
import json


# 模型处理总控
class Controller:
    def __init__(self, model_id=-1):
        self.model_id = model_id
        self.state = 'none' if self.model_id == -1 else select(f"select state from model where model_id = '{self.model_id}'")[0]['state']

    def process(self, state, input_json={}):
        if state == 'upload':  # 上传阶段
            if input_json == {}:
                INFO("ERROR", f"ModelManage Error: Upload Model need message!")
            else:
                checker = Checker()
                # 注册模型
                checker.register(input_json)
                self.model_id = input_json[0]['model_id']
                # 状态改变
                self.status('validate')
        elif state == 'validate':
            checker = Checker()
            # 检查模型 -> 推送模型(raw)
            checker.check()
            fs = ModelFS(self.model_id)
            fs.push_model('validate')
            # 状态改变
            self.status('compress')

        elif state == 'compress':
            compressor = Compressor()
            input_sub_dir = self.model_id + '/' + 'input'
            output_sub_dir = self.model_id + '/' + 'output'
            # 获取模型格式
            json_string = select(f"select format from model where model_id = '{self.model_id}'")[0]['format']
            format_json = json.loads(json_string.replace("'", '"'))
            # 拉取模型 -> 压缩 -> 推送模型(ready)
            fs = ModelFS(self.model_id)
            fs.pull_model('compress', input_sub_dir)
            model_input = os.path.join(get_path(os.path.abspath(__file__), 2, get_config("model", "compress_path")), input_sub_dir + '/' + self.model_id + "." + format_json['raw'])
            model_output = os.path.join(get_path(os.path.abspath(__file__), 2, get_config("model", "compress_path")), output_sub_dir + '/' + self.model_id + "." + format_json['ready'])
            compress_result = compressor.compress(model_input=model_input,
                                                  model_output=model_output,
                                                  config=input_json)
            # 更新结果-开始
            json_string = select(f"select json from model where model_id = '{self.model_id}'")[0]['json']
            if json_string == '':
                json_string = "{}"
            json_string = json_string.replace("None", '{}')
            other_json = json.loads(json_string.replace("'", '"'))
            other_json["compress"] = compress_result
            update({'json': other_json}, f'model_id = "{self.model_id}"', 'model')
            # 更新结果-结束
            fs.push_model('compress', output_sub_dir)
            # 状态改变
            self.status('evaluate')

        elif state == 'evaluate':
            # 获取模型格式
            json_string = select(f"select format from model where model_id = '{self.model_id}'")[0]['format']
            format_json = json.loads(json_string.replace("'", '"'))
            # 拉取模型 -> 测试模型 -> 推送模型测试结果
            fs = ModelFS(self.model_id)
            fs.pull_model('evaluate')
            # 初始化
            model_input = os.path.join(get_path(os.path.abspath(__file__), 2, get_config("model", "evaluate_path")), self.model_id + "." + format_json['ready'])
            test_output = get_path(os.path.abspath(__file__), 2, get_config("model", "result_path"))
            appraiser = Appraiser(model_id=self.model_id,
                                  model_input=model_input,
                                  test_output=test_output,
                                  json=input_json)
            evaluate_result = appraiser.start()
            # 更新结果-开始
            json_string = select(f"select json from model where model_id = '{self.model_id}'")[0]['json']
            if json_string == '':
                json_string = "{}"
            json_string = json_string.replace("None", '{}')
            other_json = json.loads(json_string.replace("'", '"'))
            other_json["evaluate"] = evaluate_result
            update({'json': other_json}, f'model_id = "{self.model_id}"', 'model')
            # 更新结果-结束
            # 状态改变
            self.status('ready')

        elif state == 'ready':
            pass
        else:
            INFO("ERROR", f"ModelManage Error: Wrong status!")

    def status(self, state):
        self.state = state
        INFO("INFO", f"Update Model {self.model_id} {self.state}")
        update_json = {'state': self.state}
        update(update_json, f'model_id = "{self.model_id}"', 'model')


if __name__ == '__main__':
    # 调用示例
    # 上传
    # controller = Controller()
    # controller.process('upload', input_json=[{'model_id':'test3', 'name':'cs', 'location':'remote', 'type':'recognition', 'attribute':'测试', 'state':'upload', 'format':{"raw":"pth","ready":"onnx"}, 'json':''}])
    # 验证
    # controller = Controller(model_id='test3')
    # controller.process('validate')
    # 压缩
    #controller = Controller(model_id='test3')
    #controller.process('compress', {'algorithm_id':1, 'dataset_id':'imagenet_test', 'input_size': 224})
    # 测试
    # controller = Controller(model_id='test7')
    # controller.process('evaluate', {'device':'gpu', 'dataset_id':'imagenet_test'})
    pass